# Created by Massimo Di Pierro - BSD License

from psim import PSim

def mergesort(A, p=0, r=None):
    if r is None: r = len(A)
    if p<r-1:
        q = int((p+r)/2)
        mergesort(A,p,q)
        mergesort(A,q,r)
        merge(A,p,q,r)


def merge(A,p,q,r):
    B,i,j = [],p,q
    while True:
        if A[i]<=A[j]:
            B.append(A[i])
            i=i+1
        else:
            B.append(A[j])
            j=j+1
        if i==q:
            while j<r:
                B.append(A[j])
                j=j+1
            break
        if j==r:
            while i<q:
                B.append(A[i])
                i=i+1
            break
    A[p:r]=B



def mergesort_p(n,d):
    p = 2**d
    comm =  PSim(p).getComm()
    if comm.rank==0:
        A = [i**2 for i in range(n)]
        A.reverse()
    else:
        A = None
    
    ds = range(d)
    for k in ds:
        ### distribute data
        shift = 2**(d-1-k)
        senders = range(0,2**d,2**(d-k))
        receivers = [sender+shift for sender in senders]
        if comm.rank in senders:
            receiver = comm.rank+shift
            print comm.rank, 'sending', receiver
            comm.send(receiver,A[n/(2**(k+1)):n/(2**k)])
        elif comm.rank in receivers:
            sender =  comm.rank-shift
            A = comm.recv(sender)
    
    #print comm.rank, A[:n/p]
    # computation
    mergesort(A,0,n/p)

    ds.reverse()
    for k in ds:
        ### distribute data
        shift = 2**(d-1-k)
        senders = range(0,2**d,2**(d-k))
        receivers = [sender+shift for sender in senders]
        if comm.rank in senders:
            receiver = comm.rank+shift
            A[n/(2**(k+1)):n/(2**k)] = comm.recv(receiver)            
            merge(A,0,n/(2**(k+1)),n/(2**k))
        elif comm.rank in receivers:
            sender =  comm.rank-shift
            comm.send(sender, A[0:n/(2**(k+1))])
    if comm.rank==0:
        print A

mergesort_p(16,2)
